import datetime

from collections import OrderedDict

from processor.models import *
from processor.utils import *

class MACross():
    @staticmethod
    def __cross_A(yday, tday):
        return tday['ma5'] > yday['ma5'] and tday['ma10'] > yday['ma10'] and \
            yday['ma5'] <= yday['ma10'] and tday['ma5'] > tday['ma10']
    
    @staticmethod
    def __cross_B(yday, tday):
        return tday['ma5'] > yday['ma5'] and tday['ma20'] > yday['ma20'] and \
            yday['ma5'] <= yday['ma20'] and tday['ma5'] > tday['ma20']

    @staticmethod
    def __cross_C(yday, tday):
        return tday['ma10'] > yday['ma10'] and tday['ma20'] > yday['ma20'] and \
            yday['ma10'] <= yday['ma20'] and tday['ma10'] > tday['ma20']

    @staticmethod
    def __sch_level_1(yday, tday):
        ''' 一级筛选 << yday: 昨日数据， tday:今日数据
                条件：交叉A
                返回：Bool
                    - 满足： True
                    - 否则： False
        '''
        return MACross.__cross_A(yday, tday)

    @staticmethod
    def __sch_level_2(yday, tday):
        ''' 二级筛选 << yday: 昨日数据， tday:今日数据
                满足条件1: 进入池1
                满足条件2: 进入池2
                满足条件3: 进入池3
                不满足: 退出
                    - 返回 (new_pool, new_stat)
                        new_pool: 观察池变动 （0~3）
                        new_stat: 池内时间变动 (-1~5)
        '''
        if tday['ma20'] > max(tday['ma5'], tday['ma10']): return (1, 0)
        elif tday['ma5'] > tday['ma20'] > tday['ma10'] and \
            tday['ma20'] > yday['ma20']: 
            return (2, 0)
        elif tday['ma20'] < min(tday['ma5'], tday['ma10']) and \
            tday['ma20'] > yday['ma20']: 
            return (3, -1)
        else: return (0, 0)

    @staticmethod
    def __sch_level_3(pool, stat, yday, tday):
        ''' 三级筛选 << pool:观察池，stat:池内时间，yday: 昨日数据， tday:今日数据
                对不同观察池中股票分别进行判断
                    - 返回 (new_pool, new_stat)
                        new_pool: 观察池变动 （0~3）
                        new_stat: 池内时间变动 (-1~5)
        '''
        if pool == 1:
            if MACross.__cross_B(yday, tday) and MACross.__cross_C(yday, tday):
                return (3, -1)
            elif MACross.__cross_B(yday, tday):
                return (2, stat + 1) if stat < 5 else (0, -1)
            else:
                return (1, stat + 1) if stat < 5 else (0, -1)
        elif pool == 2:
            if MACross.__cross_C(yday, tday):
                return (3, -1)
            else:
                return (2, stat + 1) if stat < 5 else (0, -1)
        elif pool == 3:
            return (0, -1)

    @classmethod
    def ts_sch_pcs(cls, start_date=None, end_date=None, callback=print):
        ''' Tushare 交叉筛选进程入口 << start_date: 开始日期, end_date: 结束日期, callback: 进度回调函数
                - 一、二、三级筛选
        '''
        cls.__sch_pcs(start_date=start_date, end_date=end_date, source=1, callback=callback)
    
    @classmethod
    def wd_sch_pcs(cls, start_date=None, end_date=None, callback=print):
        ''' Wind 交叉筛选进程入口 << start_date: 开始日期, end_date: 结束日期, callback: 进度回调函数
                - 一、二、三级筛选
        '''
        cls.__sch_pcs(start_date=start_date, end_date=end_date, source=2, callback=callback)

    @classmethod
    def __sch_pcs(cls, start_date=None, end_date=None, source=1, callback=print):
        ''' 内部方法 筛选进程入口 << start_date: 开始日期, end_date: 结束日期, source: 源, callback: 进度回调函数
                - 一、二、三级筛选
        '''
        if not start_date or not end_date or end_date < start_date:
            return Logger.warning(msg="均线交叉筛选日期参数错误，本次未作任何修改。", module='function')
            
        
        start_date = date_format(start_date, str)
        end_date = date_format(end_date, str)
        
        stock_count = StockBasic.objects.count()
        callback(0, stock_count)
        
        for prog, stock in enumerate(StockBasic.objects.iterator(chunk_size=16)):
            index = (stock.tushare_index if source == 1 else stock.wind_index)
            source_data = (stock.tushare_source if source == 1 else stock.wind_source)

            try:
                # 计算需要一天的提前量
                pre_start_date = start_date
                for d in index:
                    if start_date > str(d): pre_start_date = d
                    else: break
                # 交易日日期范围
                date_range = sorted([str(d) for d in index if pre_start_date <= str(d) <= end_date])
                if len(date_range) < 1:
                    continue
                # 获取其有效的上一个交易日状态
                pre_records_set = MACPool.objects.filter(source = source).filter(fdate = pre_start_date)
                pre_pool, pre_status = 0, 0
                if pre_records_set.exists():
                    for p, _ in MACPool.POOLS:
                        if not pre_records_set.filter(pid = p).exists():
                            continue
                        pre_record = pre_records_set.filter(pid = p).first()
                        if stock.scode in pre_record.field:
                            pre_pool, pre_status = p, pre_record.field[stock.scode]['status']
            except Exception as e:
                Logger.error(msg=f"均线交叉筛选交易日拉取失败, 股票代码：{stock.sname}, {e}", module='function')
            
            for pre, cur in zip(date_range[:-1], date_range[1:]):
                try:
                    # 一级筛选
                    if pre_pool == 0 and MACross.__sch_level_1(source_data[pre], source_data[cur]):
                        # 二级筛选
                        new_pool, new_status = MACross.__sch_level_2(source_data[pre], source_data[cur])
                        if not new_pool: continue
                        # 保存记录
                        MACPool.inject(
                            pool = new_pool, 
                            source = source,
                            date = cur,
                            record = {
                                stock.scode: {
                                    'sname': stock.sname,
                                    'industry': stock.industry,
                                    'status': new_status,
                                }
                            }
                        )
                        pre_pool, pre_status = new_pool, new_status
                    
                    # 三级筛选
                    elif pre_pool != 0:
                        new_pool, new_status = MACross.__sch_level_3(pre_pool, pre_status,source_data[pre], source_data[cur], )
                        if not new_pool: continue
                        # 保存记录
                        MACPool.inject(
                            pool = new_pool, 
                            source = source,
                            date = cur,
                            record = {
                                stock.scode: {
                                    'sname': stock.sname,
                                    'industry': stock.industry,
                                    'status': new_status,
                                }
                            }
                        )
                        pre_pool, pre_status = new_pool, new_status
                except Exception as e:
                    Logger.error(msg=f"均线交叉筛选/记录保存失败, 股票代码：{stock.sname}, 筛选日期：{pre}~{cur}, {e}", module='function')

            callback(prog, stock_count)


class RSIFilter():
    BATCH_SIZE = 32

    @classmethod
    def ext_pcs(cls, callback=print):
        ''' 极值监控入口 << start_date: 开始日期, end_date: 结束日期, callback: 回调函数
                全局参数 EXTREME_MONITOR_INTERVAL: 范围极值监控区间
        '''
        # if not start_date or not end_date or end_date < start_date:
        #    Logger.warning(msg="极值监控筛选日期参数错误，本次未作任何修改。", module='function')
        #    return

        stock_count = StockRSI.objects.count()
        callback(0, stock_count)

        instances = []

        for prog, stock in enumerate(StockRSI.objects.values(
            'id', 'rsi', 'wrsi', 'scode', 'stock_basic__sname', 'stock_basic__industry'
        ).iterator(chunk_size=cls.BATCH_SIZE)):
            for period in [1, 2]:
                if period == 1: rsi = stock.pop('rsi')
                else: rsi = stock.pop('wrsi')

                if not rsi:
                    continue
                    
                extre, _ = RSIExtremum.objects.get_or_create(
                    period=period,
                    scode=stock['scode'],
                    defaults=dict(
                        stock_rsi_id=stock['id'], 
                        sname=stock['stock_basic__sname'], 
                        industry=stock['stock_basic__industry'],
                    )
                )

                # 历史数据列表
                hist = []
                # 周期数据字典
                interval = OrderedDict()

                for d, v in sorted(rsi.items()):
                    date = date_format(d, datetime.date)

                    extre.updated_field = 0
                    
                    # 比较历史数据
                    if len(hist):
                        if v > max(hist):
                            extre.set_extremum('rsi_maxd', v, date)
                        if v < min(hist):
                            extre.set_extremum('rsi_mind', v, date)

                    # 比较周期数据
                    if len(interval):
                        if v > max(interval.values()):
                            extre.set_extremum('rsi_maxdx', v, date)
                        if v < min(interval.values()):
                            extre.set_extremum('rsi_mindx', v, date)

                    # 更新当日数据
                    hist.append(v)
                    interval[d] = v
                    extre.lastest_rsi = v

                    # 弹出超出范围的周期数据
                    interval_start_date = date - datetime.timedelta(days=EXTREME_MONITOR_INTERVAL)
                    for k in list(interval.keys()):
                        if date_format(k, datetime.date) <= interval_start_date:
                            interval.pop(k)
                        else:
                            break
                
                instances.append(extre)

            callback(prog, stock_count)

        # 最新日期
        latest_date = max(
            getattr(extre, field)
            for extre in instances
            for field in ['rsi_maxd_date', 'rsi_mind_date', 'rsi_maxdx_date', 'rsi_mindx_date']
        )
        
        # 非最新日期的极值 不登记更新字段
        for extre in instances:
            for field in ['rsi_maxd_date', 'rsi_mind_date', 'rsi_maxdx_date', 'rsi_mindx_date']:
                if getattr(extre, field) == latest_date:
                    break
            else:
                extre.updated_field = 0
        
        # 批量保存
        RSIExtremum.objects.bulk_update(
            instances, 
            fields=[
                'stock_rsi', 'lastest_rsi', 'rsi_maxd', 'rsi_mind', 'rsi_maxdx', 
                'rsi_mindx', 'rsi_maxd_date', 'rsi_mind_date', 'rsi_maxdx_date', 
                'rsi_mindx_date', 'updated_field', 'last_update'
            ],
            batch_size=cls.BATCH_SIZE
        )


    @classmethod
    def dly_pcs(cls, callback=print):
        ''' 每日数据分析入口 << callback: 回调函数
        '''
        # 数据分析对象 初始化函数
        RSI_STRUCT = {}
        def get_rsi_struct(date):
            if date not in RSI_STRUCT:
                RSIStructure.objects.filter(sdate=date).delete()
                RSI_STRUCT[date] = RSIStructure(sdate=date)

            return RSI_STRUCT[date]
        
        # 进度回调
        stock_count = StockRSI.objects.count()
        callback(0, stock_count)

        # 评分函数
        scoring = lambda x: 'rb1' if x <= 10 else 'rb2' if x <= 20 else \
            'rb3' if x <= 50 else 'rb4' if x <= 80 else 'rb5' if x <= 90 else 'rb6'

        for prog, stock in enumerate(StockRSI.objects.all().only('rsi', 'scode')):
            try:
                for date, value in stock.rsi.items():
                    # 登记至对应列表 同时计数
                    struct = get_rsi_struct(date)
                    score = scoring(value)
                    getattr(struct, score).append(stock.scode)
                    setattr(struct, f'{score}_cnt', len(getattr(struct, score)))

            except Exception as e:
                Logger.error(msg=f"每日数据分析执行失败，股票代码：{stock.scode}, {e}", module='function')

            callback(prog, stock_count)

        for struct in RSI_STRUCT.values():
            struct.percentage = [getattr(struct, f'rb{i}_cnt') for i in range(1, 7)]

        # 批量保存数据分析
        RSIStructure.objects.bulk_create(
            list(RSI_STRUCT.values()), batch_size=cls.BATCH_SIZE
        )

    
    @classmethod
    def prd_pcs(cls, callback=print):
        ''' 趋势判断筛选入口 callback: 回调函数
        '''
        stock_count = StockRSI.objects.count()
        callback(0, stock_count)

        rsi_trd_inst_dict = {}

        RSITrend.objects.all().delete()

        # 抓取指定记录
        def get_rsi_trd_inst(period, date):
            if (period, date) not in rsi_trd_inst_dict:
                rsi_trd_inst_dict[(period, date)] = RSITrend(
                    fdate=date, 
                    period=period,
                )

            return rsi_trd_inst_dict[(period, date)]

        # 今日日期
        today = datetime.date.today()

        # 遍历每只股票的RSI数据, 并预抓取basic信息
        for prog, basic_id in enumerate(
            StockBasic.objects.values_list('id', flat=True).iterator(chunk_size=cls.BATCH_SIZE)
        ):
            basic = StockBasic.objects.get(id=basic_id)
            stock = StockRSI.objects.filter(stock_basic=basic).values(
                'rsi', 'rsi_index', 'wrsi', 'wrsi_index', 'scode', 
            ).first()
            if stock is None: continue

            for period in (1, 2):
                # 上市日期
                list_date = basic.list_date

                '''
                # 日级筛选 period 为 1, 周级筛选 period 为 2
                if period == 1:
                    rsi, index = stock['rsi'], stock['rsi_index']
                    # 日级筛选 上市天数(120, 450] 取 [上市日期+90, ∞)
                    if 120 < date_diff(today, list_date) <= 450:
                        index = [i for i in index if i >= str(list_date + datetime.timedelta(days=90))]
                    # 日级筛选 上市天数(450, ∞) 取 [上市日期+365, ∞)
                    elif 450 < date_diff(today, list_date):
                        index = [i for i in index if i >= str(list_date + datetime.timedelta(days=365))]
                    else:
                        continue
                else:
                    rsi, index = stock['wrsi'], stock['wrsi_index']
                    # 周级筛选 上市天数(365, 575] 取 [上市日期+210, ∞)
                    if 365 < date_diff(today, list_date) <= 575:
                        index = [i for i in index if i >= str(list_date + datetime.timedelta(days=210))]
                    # 周级筛选 上市天数(575, ∞) 取 [上市日期+365, ∞)
                    elif 575 < date_diff(today, list_date):
                        index = [i for i in index if i >= str(list_date + datetime.timedelta(days=365))]
                    else:
                        continue
                '''
                
                if period == 1:
                    rsi, index = stock['rsi'], stock['rsi_index']
                else:
                    rsi, index = stock['wrsi'], stock['wrsi_index']

                # 取 [上市日期+60, ∞)
                index = [i for i in index if i >= str(list_date + datetime.timedelta(days=60))]

                try:
                    inflect_points = {}
                    # 同时在范围中遍历3个连续日期
                    for d1, d2, d3 in zip(
                        index[:-2],
                        index[1:-1],
                        index[2:],
                    ):
                        # 数据库中RSI空值处理
                        if rsi[d1] is None or rsi[d2] is None or rsi[d3] is None:
                            continue
                        # 判断点 d2 是否满足拐点要求
                        if rsi[d1] >= rsi[d2] and rsi[d3] > rsi[d2]:
                            # 初始化拐点
                            inflect_points[d2] = {
                                'date': d2, 
                                'rsi': rsi[d2],
                                'close': basic.tushare_source.get(d2, {}).get('close', 0),
                            }
                    
                    # 拐点索引日期列表
                    inflect_index = sorted(inflect_points.keys())
                    # 将索引写入拐点
                    for point_index, point_date in enumerate(inflect_index):
                        inflect_points[point_date]['index'] = point_index

                    # 初始化 M 点
                    M = None

                    for current_date in inflect_index:
                        # 更新 N 点
                        N = inflect_points[current_date]

                        # 更新 N-1 点
                        N1 = inflect_points[inflect_index[N['index'] - 1]] if N['index'] else None

                        # 更新 M 点
                        if M is None or N['rsi'] < M['rsi']:
                            M = N

                        # 更新 M 区间 更新 MC, K 点
                        if N['index'] > 2:
                            M1_TO_2N = [inflect_points[fwd] for fwd in inflect_index[M['index'] + 1: N['index']]]
                            MC = min(M1_TO_2N, key=lambda x: x['close'], default=None)
                            K = min(M1_TO_2N, key=lambda x: x['rsi'], default=None)
                        else:
                            MC = K = None

                        # 更新 K 区间 更新 KC, P 点
                        if K is not None:
                            K1_TO_2N = [inflect_points[fwd] for fwd in inflect_index[K['index'] + 1: N['index']]]
                            KC = min(K1_TO_2N, key=lambda x: x['close'], default=None)
                            P = min(K1_TO_2N, key=lambda x: x['rsi'], default=None)
                        else:
                            KC = P = None

                        # 更新 P 区间 更新 PC 点
                        if P is not None:
                            P1_TO_2N = [inflect_points[fwd] for fwd in inflect_index[P['index'] + 1: N['index']]]
                            PC = min(P1_TO_2N, key=lambda x: x['close'], default=None)
                            V = min(P1_TO_2N, key=lambda x: x['rsi'], default=None)
                        else:
                            PC = V = None

                        # 更新 V 区间 更新 VC 点
                        if V is not None:
                            V1_TO_2N = [inflect_points[fwd] for fwd in inflect_index[V['index'] + 1: N['index']]]
                            VC = min(V1_TO_2N, key=lambda x: x['close'], default=None)
                            W = min(V1_TO_2N, key=lambda x: x['rsi'], default=None)
                        else:
                            VC = W = None

                        # 更新 W 区间 更新 WC 点
                        if W is not None:
                            W1_TO_2N = [inflect_points[fwd] for fwd in inflect_index[W['index'] + 1: N['index']]]
                            WC = min(W1_TO_2N, key=lambda x: x['close'], default=None)
                        else:
                            WC = None
                        
                        # 初始化 当前组合
                        scheme = {
                            'scode': basic.scode,
                            'sname': basic.sname,
                            'industry': basic.industry,
                            'point_n': N,
                        }
                        
                        # 获取当前要写入的数据库记录
                        inst = get_rsi_trd_inst(period, current_date)

                        # 判断是否 该点的RSI小于N点 且CLOSE大于N点
                        satisfy = lambda X, N: N['rsi'] > X['rsi'] and N['close'] < X['close']

                        # 执行 Dbg-A 筛选
                        if MC is not None and satisfy(M, N) and N['close'] < MC['close']:
                            inst.bg_a.append({**scheme, 'check_point': 'm', 'point_m': M})
                        # 执行 Dbg-B 筛选
                        if KC is not None and satisfy(K, N) and N['close'] < KC['close']:
                            inst.bg_b.append({**scheme, 'check_point': 'k', 'point_k': K})
                        # 执行 Dbg-C 筛选
                        if PC is not None and satisfy(P, N) and N['close'] < PC['close']:
                            inst.bg_c.append({**scheme, 'check_point': 'p', 'point_p': P})
                        # 执行 Dbg-D 筛选
                        if VC is not None and satisfy(V, N) and N['close'] < VC['close']:
                            inst.bg_d.append({**scheme, 'check_point': 'v', 'point_v': V})
                        # 执行 Dbg-E 筛选
                        if WC is not None and satisfy(W, N) and N['close'] < WC['close']:
                            inst.bg_e.append({**scheme, 'check_point': 'w', 'point_w': W})
                        # 执行 Dbg-F 筛选
                        if N1 is not None and satisfy(N1, N):
                            inst.bg_f.append({**scheme, 'check_point': 'n1', 'point_n1': N1})

                except Exception as e:
                    Logger.error(msg=f"趋势判断筛选执行失败，股票代码：{basic.scode}, {e}", module='function')

            callback(prog, stock_count)

        # 去除全空数据
        inst_list = [inst for inst in rsi_trd_inst_dict.values() if any([inst.bg_a, inst.bg_b, inst.bg_c, inst.bg_d, inst.bg_e])]
        # 批量保存
        RSITrend.objects.bulk_create(inst_list, batch_size=cls.BATCH_SIZE)
        